# 实验管理
from pprint import pprint

import pandas as pd
from qlib.backtest import executor, backtest
from qlib.contrib.strategy import TopkDropoutStrategy
from qlib.data.dataset import DatasetH, DataHandlerLP
from qlib.utils.time import Freq
from qlib.workflow import R
from qlib.workflow.record_temp import SignalRecord, SigAnaRecord, PortAnaRecord


def get_strategy(pred_score, topK=50, n_drop=5):
    pred_score = pred_score.dropna()
    STRATEGY_CONFIG = {
        "topk": 50,
        "n_drop": 5,
        "signal": pred_score  # pred_score,
    }
    strategy_obj = TopkDropoutStrategy(**STRATEGY_CONFIG)
    return strategy_obj


def get_executor():
    EXECUTOR_CONFIG = {
        "time_per_step": "day",
        "generate_portfolio_metrics": True,
    }
    executor_obj = executor.SimulatorExecutor(**EXECUTOR_CONFIG)
    return executor_obj


def do_backtest(executor, strategy, benchmark='SH000300', start_time="2022-07-01", end_time="2022-07-31"):
    backtest_config = {
        "start_time": start_time,
        "end_time": end_time,
        "account": 100000000,
        "benchmark": benchmark,
        "exchange_kwargs": {
            "freq": 'day',
            "limit_threshold": 0.095,
            "deal_price": "close",
            "open_cost": 0.0005,
            "close_cost": 0.0015,
            "min_cost": 5,
        },
    }
    portfolio_metric_dict, indicator_dict = backtest(executor=executor, strategy=strategy, **backtest_config)
    analysis_freq = "{0}{1}".format(*Freq.parse('day'))
    print(analysis_freq)
    pprint(portfolio_metric_dict)
    pprint(indicator_dict)
    # backtest info
    report_normal_df, positions_normal = portfolio_metric_dict.get(analysis_freq)
    pprint(report_normal_df)
    pprint(positions_normal)
    return report_normal_df, positions_normal



class ExpMgr:
    def visual(self, report_normal):
        import plotly
        from qlib.contrib.report.analysis_position.report import report_graph
        from qlib.contrib.report.analysis_position import score_ic_graph, score_ic
        fig = report_graph(report_normal, show_notebook=False)
        plotly.offline.plot(fig[0], filename="result.html")

    def run_simple(self, model, dataset: DatasetH):
        model.fit(dataset)
        pred = model.predict(dataset)
        if isinstance(pred, pd.Series):
            print('is serice...')
            #pred.name = 'score'
            pred = pred.to_frame('score')

        #pprint(pred)
        #pred.to_pickle('output/pred.pkl')

        # 取label并保存
        label = dataset.prepare(segments="test", col_set="label", data_key=DataHandlerLP.DK_R)
        pprint(label)

        all = pd.concat([pred, label], axis=1)
        pprint(all)

        ic = all.groupby('datetime').apply(lambda df: df["score"].corr(df["LABEL0"]))
        ric = all.groupby('datetime').apply(lambda df: df["score"].corr(df["LABEL0"], method="spearman"))
        # label.to_pickle('output/label.pkl')
        pprint(ic)
        pprint(ric)

        group = all.groupby(level='datetime')

        def N(x):
            return int(len(x) * 0.2)

        # 打分的前20%， 与后20% 对应的 真实收益率
        r_long = group.apply(lambda x: x.nlargest(N(x), columns="score").LABEL0.mean())
        r_short = group.apply(lambda x: x.nsmallest(N(x), columns="score").LABEL0.mean())
        r_avg = group.LABEL0.mean()

        r_long_short = (r_long - r_short) / 2
        print((r_long - r_short) / 2, r_avg)

        print("多空年化收益", r_long_short.mean() * 252)
        print("多空夏普比", r_long_short.mean() / r_long_short.std() * 252 ** 0.5)
        print("平均年化收益", r_avg.mean() * 252)
        print("平均年化夏普比", r_avg.mean() / r_avg.std() * 252 ** 0.5)

        strategy_obj = get_strategy(pred_score=pred.dropna())
        executor_obj = get_executor()
        report_normal_df, positions_normal = do_backtest(executor=executor_obj, strategy=strategy_obj)
        self.visual(report_normal_df)

    def run(self, model, dataset):
        # start exp
        with R.start(experiment_name="workflow"):
            # train
            # R.log_params(**flatten_dict(task))
            model.fit(dataset)

            # prediction
            recorder = R.get_recorder()
            sr = SignalRecord(model, dataset, recorder)
            sr.generate()

            # Signal Analysis
            sar = SigAnaRecord(recorder)
            sar.generate()

            # backtest. If users want to use backtest based on their own prediction,
            # please refer to https://qlib.readthedocs.io/en/latest/component/recorder.html#record-template.

            port_analysis_config = {
                "executor": {
                    "class": "SimulatorExecutor",
                    "module_path": "qlib.backtest.executor",
                    "kwargs": {
                        "time_per_step": "day",
                        "generate_portfolio_metrics": True,
                    },
                },
                "strategy": {
                    "class": "TopkDropoutStrategy",
                    "module_path": "qlib.contrib.strategy.signal_strategy",
                    "kwargs": {
                        "signal": (model, dataset),
                        "topk": 50,
                        "n_drop": 5,
                    },
                },
                "backtest": {
                    "start_time": "2017-01-03",
                    "end_time": "2022-08-01",
                    "account": 100000000,
                    "benchmark": 'SH000300',
                    "exchange_kwargs": {
                        "freq": "day",
                        "limit_threshold": 0.095,
                        "deal_price": "close",
                        "open_cost": 0.0005,
                        "close_cost": 0.0015,
                        "min_cost": 5,
                    },
                },
            }
            par = PortAnaRecord(recorder, port_analysis_config, "day")
            par.generate()
